"""
Extract SSH from raw HYCOM bindary files and save them as nc (netcdf) files


Author: Badarvada Yadidya
Created on: May 2023

"""


import os
import os.path as op
import sys
import pandas as pd
import numpy as np
import xarray as xr

# https://github.com/abozec/BB86_PACKAGE.git
pdir = op.join('/home1/ybadarva/BB86_PACKAGE/PYTHON/')
sys.path.insert(0, pdir)

from hycom.info import read_field_names, read_field_grid_names
from hycom.io import read_hycom_fields, read_hycom_grid, sub_var2

import calendar
import argparse

parser = argparse.ArgumentParser(description="Calculate the day of the year.")
parser.add_argument("year", type=int, help="Year")
parser.add_argument("month", type=int, help="Month")

args = parser.parse_args()

year = args.year
input_month = args.month

ddir = f"/nobackup/ybadarva/2018/{input_month}/"
# ddir = f"/nfs/turbo/lsa-arbic/yadidya/hycom/{input_month}/"
ddir2 = f"/nobackup/ybadarva/2018/{input_month}/"
total_days = 0
for month in range(1, input_month):  # Iterate through months up to (input_month - 1)
    _, days_in_month = calendar.monthrange(year, month)
    total_days += days_in_month

# Now get the number of days in the input_month
_, days_in_input_month = calendar.monthrange(year, input_month)

#for day in range(1,2,1):  # Iterate through each day of the input_month
for day in range(1, days_in_input_month + 1):  # Iterate through each day of the input_month
    day_of_year = total_days + day
    day_str = f"{day_of_year:03d}"
    
    for hour in range(24):  # Iterate through each hour of the day
        hour_str = f"{hour:02d}"
        filename = f"216_archs.{year}_{day_str}_{hour_str}.a"
        print(filename)
        file_data = op.join(ddir+filename)
        hycom_field = read_hycom_fields(file_data, ['srfhgt', 'steric'])
        date = pd.DatetimeIndex([pd.Timestamp(f"{year}-{input_month}-{day} {hour_str}:00")])
        print(date)
        dt1 = xr.DataArray(hycom_field['srfhgt'],dims=['time','y', 'x'], coords={'time': date}, name='srfhgt')
        dt2 = xr.DataArray(hycom_field['steric'],dims=['time','y', 'x'], coords={'time': date}, name='steric')
        dt3 = xr.merge([dt1,dt2])
        
        dt3.to_netcdf(ddir2+f"216_archs.{year}_{day_str}_{hour_str}.nc")
        print('Done')

    
